import re
import random

from Basic_Processing import *
from myseq import MySeq

def upper_dna_seq(seq:str):
    validate_dna(seq)
    seq = seq.upper()
    seq_counter = frequency(seq)
    purines = seq_counter["A"] + seq_counter["G"]
    pyrimidines = seq_counter["T"] + seq_counter["C"]
    return seq,purines,pyrimidines

def is_self_reverse_complement(seq):
    seq = MySeq(seq,'DNA')
    seq.validate()
    return str(seq.reverse_comp()) == str(seq)

def CG_duplets_count(seq:str):
    validate_dna(seq)
    return seq.upper().count("CG")

def size_of_first_protein(seq):
    size = -1
    prot_seq = translate_seq(seq,ini_pos=random.randint(0,2))
    for seq in all_proteins_rf(prot_seq):
        size = len(seq)
        break
    return size

def enable_to_protein(seq):
    return not re.findall("(M\w*)_",seq) == []


def get_aa_dict(seq):
    seq = str(MySeq(seq,"DNA").translate())
    seq.replace("_","")
    counter = Counter(seq)
    aa_num =  sum(counter.values())
    for k,v in counter.items():
        counter[k] = v/aa_num
    return counter

def search_sub_seq(dna_seq,prot_seq):
    MySeq(dna_seq,"DNA").validate()
    MySeq(prot_seq,"PROTEIN").validate()
    pattern = re.compile(f"{prot_seq}")
    res = list()
    dna_seq = (dna_seq,reverse_complement(dna_seq))
    for index,step in zip((0,1),(1,-1)):
        for ini_pos in (0,1,2):
            trans_prot = translate_seq(dna_seq[index],ini_pos)
            for sub_port in re.finditer(pattern,trans_prot):
                s,e = sub_port.span()
                s = ini_pos + 3*s; e = ini_pos + 3*e
                if step > 0:
                    res.append(dna_seq[0][s:e])
                else:
                    res.append(dna_seq[0][::-1][s:e][::-1])
    return res

def get_repeat_sub(seq,k:int) -> dict:
    sub_dict = Counter([seq[i:i+k]for i in range(0,len(seq)-k+1)])
    result = dict()
    for k,v in sub_dict.items():
        if v >= 2:
            result[k] = v
    return result
        

if __name__ == "__main__":
    print("### ex1 ###")
    print(upper_dna_seq("ATCTCTGATCGCGCTAGA"))
    print("### ex2 ###")
    print(is_self_reverse_complement("ATCGCGAT"),
          is_self_reverse_complement("ATCGGCAT"))
    print("### ex3 ###")
    print(CG_duplets_count("CGCGACTGCGC"))
    print("### ex4 ###")
    print(size_of_first_protein("TCGATCAGCATGCTTAGCTAGCAGCTGACATGCt"))
    print("### ex5 ###")
    print(enable_to_protein("SDSFMASDFA"),enable_to_protein("DSFMASD_S"))
    print("### ex6 ###")
    print(get_aa_dict("ATCTCTGATCGCGCTAGA"))
    print("### ex7 ###")
    print(search_sub_seq("ATCTCTGATCGCGCTAGA","RD"))
    print("### ex8 ###")
    print(get_repeat_sub("CGTCGCGACGCG",4))
    pass